// This thinly wraps libserial and provides our simple VFS interface.

#include <assert.h>
#include <clock/clock.h>
#include <cspace/cspace.h>
#include <fcntl.h>
#include <libut.h>
#include <sel4/sel4.h>
#include <serial/serial.h>
#include <sos.h>
#include <stdlib.h>
#include <string.h>

#include "console.h"
#include "list.h"
#include "macros.h"
#include "pagetable.h"
#include "picoro.h"
#include "process.h"
#include "syscall.h"
#include "task.h"
#include "vfs.h"
#include "vmem_layout.h"

#include "debug.h"
#define verbose SOS_DEBUG_CONSOLE
#include <sys/debug.h>
#include <sys/panic.h>

#define CONSOLE_MAX_SEND (1450)
#define CONSOLE_LINE_BUFFERING (1)

// Console maintains its own FD table
typedef struct console_fd_table {
    struct console_fd_entry *table;
    seL4_Word readers;
    seL4_Word writers;
} console_fd_table;

typedef struct console_fd_entry {
    seL4_Word fd;
    int can_read;
    int can_write;
    UT_hash_handle hh;
} console_fd_entry;

extern struct serial *serial;
static console_fd_table _fd_table;

static timestamp_t _init_ts;
static timestamp_t _last_ts;

// Only one read is allowed at a time, and its worker task is stored here.
static coro reader_task;
// Unlimited write tasks are allowed to coexist, and they are thrown onto
// the system 'ready queue' by _write as anonymous (untracked) tasks

static int _mount(sos_mount *mount, const char *fs_data);
static int _umount(sos_mount *mount, const char *fs_data);

static int _open(sos_pcb *pcb, const char* path, fmode_t mode,
                 int proc_fd, seL4_Word *fs_fd);
static int _close(sos_pcb *pcb, seL4_Word fs_fd);
static int _read(sos_pcb *pcb, seL4_Word fs_fd, char* buf, size_t nbyte);
static int _write(sos_pcb *pcb, seL4_Word fs_fd, char* buf, size_t nbyte);
static int _getdirent(sos_pcb *pcb, int pos, size_t max);
static int _stat(sos_pcb *pcb, const char *path);

/* page table callback to wake up a writer coroutine */
static void _wakeup_cb(int err, seL4_Word sos_vaddr, void *_job) {
    coro job = (coro) _job;
    resume(job, (void *) err);
    resume(job, (void *) sos_vaddr);
    return;
}

/* SOS task to wake up a writer coroutine */
static int _wakeup_task(void *_job) {
    coro job = (coro) _job;
    resume(job, NULL);
    return SOS_TASK_SUCCESS;
}

// The function that handles all console input
static void _serial_read_handler(struct serial *serial, char c);

int sos_console_init(sos_filesystem **fs) {
    int err;

    // If driver is in use, the caller is in error.
    if (serial != NULL) return SOS_CONSOLE_REINIT;
    dprintf(1, "\nInitialising libserial...\n");
    serial = serial_init();
    assert(serial != NULL);

    // Attach handler for console readers
    dprintf(1, "\tAttaching console_read handler...\n");
    err = serial_register_handler(serial, &_serial_read_handler);
    if (err) return SOS_CONSOLE_REINIT;

    // Alloc some memory for the fs object
    dprintf(1, "\tCreating console filesystem...\n");
    (*fs) = malloc(sizeof(struct sos_filesystem));
    if ((*fs) == NULL) return SOS_CONSOLE_OOM;

    // Populate it with our function pointers
    (**fs) = (sos_filesystem) { _mount, _umount };

    // Store initial timestamps
    _init_ts = time_stamp();
    _last_ts = time_stamp();

    return SOS_CONSOLE_SUCCESS;
}

// XXX
int sos_console_destroy(sos_filesystem *fs) {
    return SOS_CONSOLE_FAILURE;
}

// VFS interface here kinda assumed C had partially-applied functions. It don't.
// XXX Interface really needs to provide some fs-specific data field to distinguish mounts.
static int _mount(sos_mount *mount, const char *fs_data) {
    if (serial == NULL || mount == NULL) return SOS_VFS_INVAL;
    (*mount) = (sos_mount) {
        .open = &_open,
        .close = &_close,
        .read = &_read,
        .write = &_write,
        .getdirent = &_getdirent,
        .stat = &_stat
    };
    return SOS_VFS_SUCCESS;
}

// XXX M5. Too much to think about.
static int _umount(sos_mount *mount, const char *fs_data) {
    return SOS_VFS_INVAL;
}

/* There is precisely one file, and it's "/".
 * "/" is implicit when accessing a mountpoint.
 */
static inline int _valid_path(const char* path) {
    return (path && (strcmp(path, "/") <= 0));
}

static int _register_reader(seL4_Word fd) {
    console_fd_entry *_fd = (console_fd_entry *) fd;
    if (_fd_table.readers == 0) {
        _fd->can_read = 1;
        _fd_table.readers++;
        return SOS_VFS_SUCCESS;
    } else {
        dprintf(3, "console_register_reader: multiple reader attempt\n");
        return SOS_VFS_PERM;
    }
}

static int _unregister_reader(seL4_Word fd) {
    console_fd_entry *_fd = (console_fd_entry *) fd;
    if (_fd->can_read && _fd_table.readers) {
        _fd->can_read = 0;
        _fd_table.readers--;
        return SOS_VFS_SUCCESS;
    } else return SOS_VFS_INVAL;
}

// Anyone can be a writer - unbounded
static int _register_writer(seL4_Word fd) {
    console_fd_entry *_fd = (console_fd_entry *) fd;
    _fd->can_write = 1;
    _fd_table.writers++;
    return SOS_VFS_SUCCESS;
}

static int _unregister_writer(seL4_Word fd) {
    console_fd_entry *_fd = (console_fd_entry *) fd;
    _fd->can_write = 0;
    _fd_table.writers--;
    return SOS_VFS_SUCCESS;
}

static int _can_write(seL4_Word fd) {
    console_fd_entry *_fd = (console_fd_entry *) fd;
    return (_fd->can_write);
}

static int _can_read(seL4_Word fd) {
    console_fd_entry *_fd = (console_fd_entry *) fd;
    return (_fd->can_read);
}
static int _open(sos_pcb *pcb, const char* path, fmode_t mode,
                 int proc_fd, seL4_Word *fs_fd) {
    int err = SOS_VFS_SUCCESS;
    seL4_CPtr reply_cap;

    if (serial == NULL || path == NULL) err = SOS_VFS_INVAL;
    if (!_valid_path(path)) err = SOS_VFS_NOENT;
    if (mode != O_RDONLY && mode != O_WRONLY && mode != O_RDWR)
        err = SOS_VFS_INVAL; // Invalid mode flag - only three allowed

    if (err) {
        sos_syscall_open_reply(pcb, proc_fd, err);
        return;
    }

    dprintf(3, "console_open: Attempt to open %.5s with perms %d\n", path, mode);

    console_fd_entry *fd_entry = malloc(sizeof(struct console_fd_entry));
    if (fd_entry == NULL) {
        sos_syscall_open_reply(pcb, proc_fd, err);
        return SOS_VFS_OOM;
    }

    // We use pointers as the fs_fd. (fs_fd is completely hidden from the process)
    // ... we also can't get use-after-free'd, such attacks would be caught by VFS
    (*fs_fd) = (seL4_Word) fd_entry;
    fd_entry->fd = (*fs_fd);

    if (mode == O_RDONLY || mode == O_RDWR) {
        dprintf(5, "console_open: attaching reader\n");
        err = _register_reader(*fs_fd);
        if (err) {
            sos_syscall_open_reply(pcb, proc_fd, SOS_VFS_PERM);
            return SOS_VFS_PERM;
        }
    }

    if (mode == O_WRONLY || mode == O_RDWR) {
        dprintf(5, "console_open: attaching writer\n");
        err = _register_writer(*fs_fd);
        if (err) {
            sos_syscall_open_reply(pcb, proc_fd, SOS_VFS_PERM);
            return SOS_VFS_PERM;
        }
    }

    // Add entry to the FD table
    HASH_ADD_INT(_fd_table.table, fd, fd_entry);

    // Bump atime counter
    _last_ts = time_stamp();

    // Wake up process - we're done
    sos_syscall_open_reply(pcb, proc_fd, SOS_VFS_SUCCESS);

    return SOS_VFS_SUCCESS;
}

static int _close(sos_pcb *pcb, seL4_Word fs_fd) {
    int close_r, close_w;

    dprintf(5, "console_close: closing %p\n", fs_fd);

    // Make sure this FD exists and is in the table. If not, heap corruption...
    console_fd_entry *fde;
    HASH_FIND_INT(_fd_table.table, &fs_fd, fde);
    if (fde == NULL) return SOS_VFS_INVAL;

    // Make sure this FD can't read or write anymore
    close_r = _unregister_reader(fs_fd);
    close_w = _unregister_writer(fs_fd);

    // Should be either a reader or writer, not neither.
    // This should be enforced by _open, so this catches the absurd case.
    assert(!(close_r && close_w));

    // Remove from _fd_table
    HASH_DEL(_fd_table.table, fde);
    // Free the FD entry structure
    free(fde);

    return SOS_VFS_SUCCESS;
}

static void *_write_task(void *arg) {
    // Grab our arguments from the list on the stack
    list args = (list) arg;

    assert(args->size == 4);
    sos_pcb *pcb = (sos_pcb *) list_pop_front(args);
    char *buf = (char *) list_pop_front(args);
    size_t nbyte = (size_t) list_pop_front(args);
    coro self = (coro) list_pop_front(args);

    int err = SOS_VFS_SUCCESS;
    size_t sent = 0;
    seL4_Word cur_vaddr = (seL4_Word) buf;
    seL4_Word sos_vaddr;
    seL4_Word sos_kaddr;
    size_t max_send;
    page_attrs attrs;

    seL4_MessageInfo_t reply;

    while (sent < nbyte) {
        // Pin page, yield
        err = sos_page_pin(pcb, cur_vaddr, &_wakeup_cb, (void *) self);
        err = (int) yield(NULL);
        sos_kaddr = (seL4_Word) yield(NULL);

        // Check for abort
        if (pcb->abort && !err) {
            frame_free(sos_kaddr);
        }
        if (pcb->abort) {
            dprintf(6, "console_write: aborting operation\n");
            free(pcb);
            return NULL;
        }

        // Check if we need to destroy the process
        if (err == SOS_PAGE_TABLE_SEGFAULT) {
            dprintf(1, "PID %d segfault, access violation\n", pcb->pid);
            cspace_free_slot(cur_cspace, pcb->reply_cap);
            process_delete(pcb, NULL, NULL);
            return NULL;
        }

        // All other errors are operation-stoppers
        if (err) break;

        // Lookup, check perms
        err = sos_page_lookup_nofault(pcb, cur_vaddr, &sos_vaddr, &attrs);
        if (err || !(attrs.readable)) {
            sos_page_release(pcb, cur_vaddr);
            break;
        }

        // Make sure sos_vaddr is aligned correctly.
        sos_vaddr = PAGE_ALIGN(sos_vaddr) + (cur_vaddr & PAGE_MASK);

        // Do math. Don't want to read past a page boundary.
        // libserial also has a limit of ~1450 bytes.
        max_send = MIN3(PAGE_REMAINING(cur_vaddr),
                        CONSOLE_MAX_SEND,
                        (nbyte - sent));

        // Set up call, yield
        max_send = serial_send(serial, (char *) sos_vaddr, max_send);

        // Release page
        err = sos_page_release(pcb, cur_vaddr);
        conditional_panic(err, "Could not release pinned page");

        sent += max_send;
        cur_vaddr += max_send;

        // Don't want to yield on our final run
        if (sent < nbyte) {
            // Schedule ourselves for next tick
            err = sos_task_add_ready(&_wakeup_task, (void *) self,
                                     SOS_TASK_PRIORITY_LOW);
            if (err) break;
            yield(NULL);

            if (pcb->abort) {
                dprintf(6, "console_write: aborting operation\n");
                free(pcb);
                return NULL;
            }

        }
    }

    // Wake up the calling process
    dprintf(6, "Wrote %d bytes. Calling back...\n", sent);
    sos_syscall_write_reply(pcb, err, sent);
    return NULL;
}

static int _write(sos_pcb *pcb, seL4_Word fs_fd, char* buf, size_t nbyte) {
    int err;
    list args;
    coro job;
    console_fd_entry *fde;

    // See if any writing actually needs to be done
    if (serial == NULL) {
        sos_syscall_write_reply(pcb, SOS_VFS_INVAL, 0);
        return SOS_VFS_INVAL;
    }

    // zero-writes are not interesting
    if (nbyte == 0 || buf == NULL) {
        sos_syscall_write_reply(pcb, SOS_VFS_SUCCESS, 0);
        return SOS_VFS_SUCCESS;
    }

    // Make sure this FD exists and is in the table. If not, heap corruption...
    HASH_FIND_INT(_fd_table.table, &fs_fd, fde);
    if (fde == NULL) {
        dprintf(6, "console_write: could not find FD\n");
        sos_syscall_write_reply(pcb, SOS_VFS_INVAL, 0);
        return SOS_VFS_INVAL;
    }
    if (!_can_write(fs_fd)) {
        sos_syscall_write_reply(pcb, SOS_VFS_PERM, 0);
        return SOS_VFS_PERM;
    }

    // Create coroutine for the write
    job = coroutine(&_write_task);

    // Pass arguments to the coroutine via a list
    args = list_init();

    err  = list_push_back(args, (SOS_LIST_TYPE) pcb);
    err |= list_push_back(args, (SOS_LIST_TYPE) buf);
    err |= list_push_back(args, (SOS_LIST_TYPE) nbyte);
    err |= list_push_back(args, (SOS_LIST_TYPE) job);
    if (err) {
        // Show's over, but we should call back the process
        dprintf(3, "console_write: Could not allocate task args (%d)\n", err);
        list_destroy(args);
        sos_syscall_write_reply(pcb, SOS_VFS_OOM, 0);
        return SOS_VFS_OOM;
    }

    // Write one chunk with initial arguments
    resume(job, (void *) args);
    // Destroy the list, values have been copied into job stack
    list_destroy(args);

    return SOS_VFS_SUCCESS;
}

static void _read_pin_cb(int err, seL4_Word sos_vaddr, void *cookie) {
    coro self = (coro) cookie;
    resume(self, (void *) err);
    resume(self, (void *) sos_vaddr);
    return;
}

static void *_read_task(void *arg) {
    // Grab our arguments from the list on the stack
    list args = (list) arg;

    assert(args->size == 4);
    sos_pcb *pcb = (sos_pcb *) list_pop_front(args);
    char *buf = (char *) list_pop_front(args);
    size_t nbyte = (size_t) list_pop_front(args);
    coro self = (coro) list_pop_front(args);

    int err = SOS_VFS_SUCCESS;
    size_t sent = 0;
    seL4_Word cur_vaddr = (seL4_Word) buf;
    seL4_Word sos_vaddr;
    seL4_Word sos_kaddr;
    page_attrs attrs;
    size_t max_send;
    char c;

    seL4_MessageInfo_t reply;

    dprintf(3, "console_read: Reader task running. Want %dB\n", nbyte);

    while (sent < nbyte) {
        // Pin receiving page
        err = sos_page_pin(pcb, cur_vaddr, &_read_pin_cb, (void *) self);
        err = (int) yield(NULL);
        sos_kaddr = (seL4_Word) yield(NULL);

        if (pcb->abort && !err) {
            // REturn the pinned frame
            frame_free(sos_kaddr);
        }
        if (pcb->abort) {
            dprintf(6, "console_read: aborting operation\n");
            free(pcb);
            return NULL;
        }

        // Check if we need to destroy the process
        if (err == SOS_PAGE_TABLE_SEGFAULT) {
            dprintf(1, "PID %d segfault, access violation\n", pcb->pid);
            cspace_free_slot(cur_cspace, pcb->reply_cap);
            process_delete(pcb, NULL, NULL);
            return NULL;
        } else if (err) break;

        // Look up kernel address for pinned page
        err = sos_page_lookup_nofault(pcb, cur_vaddr, &sos_vaddr, &attrs);
        if (err) break;

        // Make sure address is properly aligned
        sos_vaddr = PAGE_ALIGN(sos_vaddr) + (cur_vaddr & PAGE_MASK);

        // Compute maximum write
        max_send = MIN(nbyte, PAGE_REMAINING(cur_vaddr));

        // Start accepting console characters
        reader_task = self;

        for (int i = 0; i < max_send; i++) {
            // Wait for a character from the handler
            arg = yield(NULL);
            // Check for abort
            if (pcb->abort) {
                dprintf(6, "console_read: aborting operation\n");
                free(pcb);
                reader_task = NULL;
                return NULL;
            }

            c = (char) (seL4_Word) arg;
            *((char *) sos_vaddr) = c;

            sos_vaddr++;
            cur_vaddr++;
            sent++;

            if (CONSOLE_LINE_BUFFERING && (c == '\n' || c == '\r')) {
                // Send back the read line by breaking out of both loops
                nbyte = sent;
                break;
            }
        }
    }

    // Wake up the calling process
    dprintf(6, "Read %d bytes. Calling back...\n", sent);
    sos_syscall_read_reply(pcb, err, sent);
    return NULL;
}

static int _read(sos_pcb *pcb, seL4_Word fs_fd, char* buf, size_t nbyte) {
    int err;
    list args;
    coro job;
    console_fd_entry *fde;

    // See if any reading actually needs to be done
    if (serial == NULL || buf == NULL) {
        sos_syscall_read_reply(pcb, SOS_VFS_INVAL, 0);
        return SOS_VFS_INVAL;
    }
    // Make sure this FD exists and is in the table. If not, heap corruption...
    HASH_FIND_INT(_fd_table.table, &fs_fd, fde);
    if (fde == NULL) {
        dprintf(3, "console_read: fs_fd %u not found in table", fs_fd);
        sos_syscall_read_reply(pcb, SOS_VFS_INVAL, 0);
        return SOS_VFS_INVAL;
    }
    if (!_can_read(fs_fd)) {
        sos_syscall_read_reply(pcb, SOS_VFS_PERM, 0);
        return SOS_VFS_PERM;
    }

    dprintf(3, "console_read: trying to start reader task\n");

    // Create coroutine for the read. This will set up the handler etc
    job = coroutine(&_read_task);

    // Pass arguments to the coroutine via a list
    args = list_init();

    err  = list_push_back(args, (SOS_LIST_TYPE) pcb);
    err |= list_push_back(args, (SOS_LIST_TYPE) buf);
    err |= list_push_back(args, (SOS_LIST_TYPE) nbyte);
    err |= list_push_back(args, (SOS_LIST_TYPE) job);
    if (err) {
        // Show's over, OOM. we should call back the process
        dprintf(3, "console_read: Could not allocate task args (%d)\n", err);
        list_destroy(args);
        sos_syscall_read_reply(pcb, SOS_VFS_OOM, 0);
        return SOS_VFS_OOM;
    }

    // Send its initial arguments
    resume(job, (void *) args);
    list_destroy(args);

    return SOS_VFS_SUCCESS;
}

static void _serial_read_handler(struct serial *serial, char c) {
    dprintf(7, "console_read: serial handler received char %c. reader task %p\n", c, reader_task);
    if (reader_task == NULL || !resumable(reader_task)) return;
    dprintf(6, "console_read: Writing %c to process memory...", c);
    resume(reader_task, (void *) (seL4_Word) c);
    dprintf(6, "done\n");
    if (!resumable(reader_task)) {
        dprintf(5, "console_read: Reader task has been detached\n");
        reader_task = NULL;
    }
}

// console is a file, not a directory, so no possible contents
static int _getdirent(sos_pcb *pcb, int pos, size_t max) {
    int err;
    err = SOS_VFS_NOENT;
    if (serial == NULL) err = SOS_VFS_INVAL;

    // Call back the process
    sos_syscall_getdirent_reply(pcb, NULL, 0, err);
    return err;
}

static int _stat(sos_pcb *pcb, const char *path) {
    int err;
    sos_stat_t stat_buf = { 0 };

    if (serial == NULL) {
        err = SOS_VFS_INVAL;
    } else if (!_valid_path(path)) {
        err = SOS_VFS_NOENT;
    } else {
        err = SOS_VFS_SUCCESS;
        stat_buf.st_type = ST_SPECIAL;
        stat_buf.st_fmode = FM_WRITE | FM_READ;
        stat_buf.st_size = 0;
        stat_buf.st_ctime = _init_ts;
        stat_buf.st_atime = _last_ts;
    }

    sos_syscall_stat_reply(pcb, err, stat_buf);
    return SOS_VFS_SUCCESS;
}
